package com.czk.polling;

import java.util.TreeMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 红黑树的平滑轮询版本
 *
 * @Author:ChenZhangKun
 * @Date: 2021/12/24 10:51
 */
public class WrrSecurityLoopTreeMap {
    private int total;
    private TreeMap<Integer, Element> pools = new TreeMap<>();
    private AtomicInteger count = new AtomicInteger();

    public WrrSecurityLoopTreeMap(Element[] elements) {
        // 计算total
        for (Element element : elements) {
            this.total += element.getWeight();
            pools.put(this.total - 1, element);
        }
    }

    public String next() {
        final int modulo = total;
        for (; ; ) {
            int hit = count.get();
            int next = (hit + 1) % modulo;
            if (count.compareAndSet(hit, next) && hit < modulo) {
                return pools.ceilingEntry(hit).getValue().getIp();
            }
        }
    }

    public static void main(String[] args) {
        Element[] elements = new Element[10];
        for (int i = 0; i < 10; i++) {
            elements[i] = new Element("192.168.25." + i, i);
        }
        WrrSecurityLoopTreeMap wrrSecurityLoopTreeMap = new WrrSecurityLoopTreeMap(elements);
        for (int i = 0; i < 10; i++) {
            System.out.println(wrrSecurityLoopTreeMap.next());
        }


    }

    static class Element {
        private String ip;
        private int weight;

        public Element(String ip, int weight) {
            this.ip = ip;
            this.weight = weight;
        }

        public String getIp() {
            return ip;
        }

        public void setIp(String ip) {
            this.ip = ip;
        }

        public int getWeight() {
            return weight;
        }

        public void setWeight(int weight) {
            this.weight = weight;
        }
    }
}
